#! /usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (c) Huoty, All rights reserved
# Author: Huoty <sudohuoty@163.com>
# CreateTime: 2020-05-22 18:04:36

"""Top K 问题

从一个数组中查找第 K 大的数
"""


def partition(array, start, end):
    """快排分区"""
    pivot = start
    pivot_value = array[pivot]  # 以第一个元素作为基准
    left = start
    right = end
    while left < right:
        while left < right and array[right] >= pivot_value:  # 右指针必须先移动
            right -= 1
        while left < right and array[left] <= pivot_value:
            left += 1
        if left < right:
            array[left], array[right] = array[right], array[left]
    if pivot != left:
        array[pivot], array[left] = array[left], array[pivot]
    return left


def quick_sort(array, start=0, end=None):
    """快速排序"""
    if end is None:
        end = len(array) - 1

    # 递归的退出条件
    if start >= end:
        return array

    pivot = partition(array, start, end)
    quick_sort(array, start, pivot - 1)
    quick_sort(array, pivot + 1, end)
    return array


def find_k(array, k):
    """快排实现"""
    arr_len = len(array)
    start = 0
    end = arr_len - 1
    povit = partition(array, start, end)
    while povit != k - 1:
        if povit < k - 1:
            povit = partition(array, povit + 1, end)
        else:
            povit = partition(array, 0, povit - 1)
    return array[povit]


if __name__ == "__main__":
    import random

    arr = [random.randint(0, 99) for _ in range(0, 20)]
    print("row array:", arr)
    print("quick sort result:", quick_sort(arr.copy()))
    # print("row array:", arr)
    print("find 3 largest number:", find_k(arr, 3))
